/te-ci pytorch

Greptile Summary

This PR improves the device-init grouped linear module by adding single grouped weight/bias support for the native grouped-tensor GEMM path, fixing numerical issues, and refactoring the capability-gate logic into two exported helper functions (is_module_grouped_tensor_path_supported / is_op_fuser_grouped_tensor_path_supported). It also fixes make_grouped_weights to no longer skip Float8CurrentScaling parameters, adds an output= reuse argument to tex.group_quantize for CUDA-graph-safe weight caching, and introduces a use_grouped_tensor constructor flag to replace the deprecated NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM environment variable.

• module/grouped_linear.py: Single grouped weight/bias is limited to the grouped-tensor path; the legacy split-quantize path only accepts discrete parameters. Weight-workspace caching, delayed wgrad, fused wgrad accumulation, and bias-gradient flow are all extended for grouped parameter layouts.
• ops/basic/grouped_linear.py: Removes the inline _is_graph_safe_path_supported method and replaces it with the exported is_op_fuser_grouped_tensor_path_supported. Single grouped parameters now raise immediately when the native path is unavailable instead of falling through to an incompatible split-quantize flow.
• C++/CUDA: group_quantize gains output= reuse support; launch_grouped_bias_add short-circuits on zero total rows; MXFP8 scale layout for uniform grouped tensors now correctly derives per-member scale bases.

Confidence Score: 4/5

PR is safe to merge; all issues found are P2 quality/maintenance concerns with no blocking defects.

The changes are well-tested with new targeted test cases covering the new single grouped weight/bias paths, CUDA graph safety, and edge cases. The key bug fixes (out_features dim, MXFP8 scale base, zero-row short-circuit, meta-device crash) are correct. The capability-gate duplication and cuBLASLt version floor change are minor P2 concerns that do not block correctness.

Files Needing Attention: module/grouped_linear.py and ops/basic/grouped_linear.py — the two is_*_grouped_tensor_path_supported predicates should ideally share logic to avoid silent divergence.

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[GroupedLinear.forward] --> B{single_grouped_weight\nor single_grouped_bias?}
    B -- Yes --> C{use_grouped_tensor?}
    C -- No --> D[RuntimeError: requires use_grouped_tensor=True]
    C -- Yes --> E{is_module_grouped_tensor_path_supported?}
    E -- No --> F[RuntimeError: native path unavailable]
    E -- Yes --> G{m_splits on CUDA?}
    G -- No --> H[ValueError: requires CUDA m_splits]
    G -- Yes --> I[_forward_grouped_tensor]
    B -- No --> J{use_grouped_tensor=True AND supported?}
    J -- Yes --> I
    J -- No --> K[Legacy split-quantize path]
    I --> L{single_grouped_weight?}
    L -- Yes --> M[_prepare_weights_for_grouped_tensor_gemm\nGroupedTensorStorage weight]
    L -- No --> N[_prepare_weights_for_grouped_tensor_gemm\nDiscrete weights list]
    M --> O{workspace cached?}
    O -- No --> P[tex.group_quantize creates workspace]
    O -- skip_fp8_weight_update --> Q[tex.group_quantize with noop_flag]
    O -- first_microbatch --> R[tex.group_quantize updates workspace]
    O -- not first_microbatch --> S[Reuse cached workspace]
    P --> T[grouped GEMM forward]
    Q --> T
    R --> T
    S --> T
    N --> T

_{Reviews (22): Last reviewed commit: "mcore integration fixes" | Re-trigger Greptile}

/te-ci pytorch

/te-ci pytorch